#include <armadillo>

#include "helper/helper.h"
#include "helper/StopWatch.h"
extern po::variables_map vm;

void zero_threshold(const arma::sp_mat& sp_A, arma::sp_mat& A)
{
  A.zeros(sp_A.n_rows, sp_A.n_cols);
  for (unsigned int j = 0; j < sp_A.n_cols; j++)
    for (unsigned int i = 0; i < sp_A.n_rows; i++)
      if (sp_A(i,j) > 1e-16) A(i,j) = sp_A(i,j);
}

void normalize_column(const arma::sp_mat& sp_A, arma::sp_mat& A)
{
  A = sp_A;
  for (unsigned int j = 0; j < A.n_cols; j++) {
    double s = arma::norm(sp_A.col(j), 2);
    if (s > 1e-16)
      A.col(j) /= s;
  }
}


void convert_sparse_to_dense(arma::sp_mat& sp_A, arma::mat& A)
{
  A = arma::mat(sp_A.n_rows, sp_A.n_cols);
  for (unsigned int i = 0; i < A.n_rows; i++)
    for (unsigned int j = 0; j < A.n_cols; j++)
      A(i,j) = sp_A(i,j);
}

void solve_sparse_anls(const arma::sp_mat& A, arma::sp_mat& W, arma::sp_mat& H)
{
  StopWatch watch;
  using namespace arma;
  int m = A.n_rows, n = A.n_cols, k = W.n_cols;
  sp_mat U(m, k), V(n, k), X(n, k), V_tmp(n, k);
  sp_mat Lambda(m, k), Pi(n, k), Delta(n, k);

  PO_REQUIRE(vm, "iter");

  PO_VARIABLE(rho, vm, "rho", double); // ADMM parameter
  PO_VARIABLE(alpha, vm, "rate", double); // dual ascent step
  PO_VARIABLE(maxiter, vm, "iter", int);

  for (int iter = 0; iter < maxiter; iter++) {
    // update W
    sp_mat HTHpI_tmp = trans(H)*H+rho*speye(k,k);
    mat HTHpI;
    convert_sparse_to_dense(HTHpI_tmp, HTHpI);
    HTHpI = arma::inv(HTHpI);
    W = (A*H-Lambda+rho*U)*HTHpI;

    // update H
    sp_mat WTWpI_tmp = trans(W)*W+rho*speye(k,k);
    mat WTWpI;
    convert_sparse_to_dense(WTWpI_tmp, WTWpI);
    WTWpI = arma::inv(WTWpI);
    H = (trans(A)*W-Pi-Delta+rho*V+rho*X)*WTWpI;

    // update U, V, X
    zero_threshold(W + Lambda / rho, U);
    zero_threshold(H + Pi / rho, V);
    //if (vm.count("normalize")) normalize_column(V_tmp, V);
    if (vm.count("normalize")) normalize_column(H + Delta / rho, X);

    // update Lambda, Pi, Delta (dual variables)
    double rate = alpha/sqrt(iter+1)*rho;
    Lambda += rate*(W-U);
    Pi += rate*(H-V);
    if (vm.count("normalize")) Delta += rate*(H-X);

    if (iter % 100 == 0 || iter == maxiter-1)
      MSG("it=" << iter << " |res|=" << norm(A-U*trans(V),"fro")  << " norm(H,fro)=" << norm(V, "fro")
          << " d=" << watch.GetDuration());
  }
  MSG(watch.GetTotalDuration());
  W = U;
  H = V;
}

void run_sparse_anls()
{
  PO_VARIABLE(m, vm, "m", int);
  PO_VARIABLE(n, vm, "n", int);
  PO_VARIABLE(k, vm, "k", int);
  PO_VARIABLE(density, vm, "density", double);

  using namespace arma;
  sp_mat A, W, H;
  W.sprandu(m, k, density);
  H.sprandu(n, k, density);
  A = W*trans(H);
  MSG_VERBOSE(vm, "A=\n" << A << "W=\n" << W << "H=\n" << H);

  srand((unsigned int)time(NULL));
  sp_mat W0, H0_tmp, H0;
  W0.sprandu(m, k, density);
  H0_tmp.sprandu(n, k, density);
  normalize_column(H0_tmp, H0);
  solve_sparse_anls(A, W0, H0);

  MSG_VERBOSE(vm, "W0=\n" << W0 << "H0=\n" << H0);
  MSG("|res|=" << norm(A-W0*trans(H0), "fro"));
  MSG("sqrt(k)=" << sqrt(k) << " norm(H,fro)=" << norm(H0, "fro"));
}
